import cv2
from skimage import util
import math
import numpy as np
import matplotlib.pyplot as plt


# 仿真运动模糊
def motion_process(image_size, motion_angle):
    PSF = np.zeros(image_size)
    center_position = (image_size[0] - 1) / 2

    slope_tan = math.tan(motion_angle * math.pi / 180)
    slope_cot = 1 / slope_tan
    if slope_tan <= 1:
        for i in range(15):
            offset = round(i * slope_tan)
            PSF[int(center_position + offset), int(center_position - offset)] = 1
        return PSF / PSF.sum()  # 对点扩散函数进行归一化亮度
    else:
        for i in range(15):
            offset = round(i * slope_cot)
            PSF[int(center_position - offset), int(center_position + offset)] = 1
        return PSF / PSF.sum()


# 对图片进行运动模糊
def make_blurred(input, PSF, eps):
    input_fft = np.fft.fft2(input)  # 进行二维数组的傅里叶变换
    PSF_fft = np.fft.fft2(PSF) + eps
    blurred = np.fft.ifft2(input_fft * PSF_fft)
    blurred = np.abs(np.fft.fftshift(blurred))
    return blurred


def inverse(input, PSF, eps):  # 逆滤波
    input_fft = np.fft.fft2(input)
    PSF_fft = np.fft.fft2(PSF) + eps  # 噪声功率，这是已知的，考虑epsilon
    result = np.fft.ifft2(input_fft / PSF_fft)  # 计算F(u,v)的傅里叶反变换
    result = np.abs(np.fft.fftshift(result))
    return result


def wiener(input, PSF, eps, K=0.01):  # 维纳滤波，K=0.01
    input_fft = np.fft.fft2(input)
    PSF_fft = np.fft.fft2(PSF) + eps
    PSF_fft_1 = np.conj(PSF_fft) / (np.abs(PSF_fft) ** 2 + K)
    result = np.fft.ifft2(input_fft * PSF_fft_1)
    result = np.abs(np.fft.fftshift(result))
    return result


if __name__ == '__main__':
    image = cv2.imread(r"C:\Users\Public\opencv\Figure\lena.jpg")
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 进行运动模糊处理
    img_h, img_w = image.shape[0:2]
    PSF = motion_process((img_h, img_w), 60)
    blurred = np.abs(make_blurred(image, PSF, 1e-3))

    plt.subplot(241), plt.axis('off')
    plt.title("1")
    plt.imshow(blurred)

    result_inv = inverse(blurred, PSF, 1e-3)  # 逆滤波
    plt.subplot(242), plt.axis('off')
    plt.title("2")
    plt.imsave(r'C:\Users\Public\opencv\Figure\lenaInv.jpg', result_inv)
    plt.imshow(result_inv)

    # 添加噪声,standard_normal产生随机的函数
    gs_noisy = util.random_noise(image, mode='gaussian')
    plt.subplot(243), plt.axis('off')
    plt.title("3")
    plt.imshow(gs_noisy)  # 显示添加噪声且运动模糊的图像

    # 对添加噪声的图像进行逆滤波
    result_invn = inverse(gs_noisy, PSF, 0.1 + 1e-3)
    plt.subplot(244), plt.axis('off')
    plt.title("4")
    plt.imsave(r'C:\Users\Public\opencv\Figure\lenaInvn.jpg', result_invn)
    plt.imshow(result_invn)
    plt.show()

    # 进行运动模糊处理
    img_h, img_w = image.shape[0:2]
    PSF = motion_process((img_h, img_w), 60)
    blurred = np.abs(make_blurred(image, PSF, 1e-3))
    plt.subplot(245), plt.axis('off')
    plt.title("5")
    plt.imshow(blurred)
    plt.imsave(r'C:\Users\Public\opencv\Figure\lenaMb.jpg', blurred)

    resultwd = wiener(blurred, PSF, 1e-3)  # 维纳滤波
    plt.subplot(246), plt.axis('off')
    plt.title("6")
    plt.imshow(resultwd)
    plt.imsave(r'C:\Users\Public\opencv\Figure\lenaMd.jpg', resultwd)

    # 添加噪声,standard_normal产生随机的函数
    blurred_noisy = blurred + 0.1 * blurred.std() * np.random.standard_normal(blurred.shape)
    plt.subplot(247), plt.axis('off')
    plt.title("7")
    plt.imshow(blurred_noisy)  # 显示添加噪声且运动模糊的图像

    # 对添加噪声的图像进行维纳滤波
    resultwdn = wiener(blurred_noisy, PSF, 0.1 + 1e-3)
    plt.subplot(248), plt.axis('off')
    plt.title("8")
    plt.imshow(resultwdn)
    plt.show()
    plt.imsave(r'C:\Users\Public\opencv\Figure\lenaWdn.jpg', resultwdn)



